#include <Eigen/Eigen>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/TwistStamped.h>
#include "ros/ros.h"
#include "nav_msgs/Path.h"
#include <iostream>
#include "tube_planning/tube_planning.h"
#include <ros/time.h>
#include "tube_planning/mmpc.h"
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <mavros_msgs/PositionTarget.h>

mavros_msgs::State current_state;
geometry_msgs::PoseStamped local_pos, ref_pose;
geometry_msgs::TwistStamped local_vel;
ros::Subscriber state_sub, local_pos_sub, local_vel_sub, trigger_sub;
ros::Publisher local_pos_pub, local_pva_pub, predict_pos_pub, ref_pos_pub;
ros::ServiceClient arming_client, set_mode_client;
int begin_task = 0;
void state_cb(const mavros_msgs::State::ConstPtr& msg){
    current_state = *msg;
}
void local_pos_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
    local_pos = *msg;
}
void local_vel_cb(const geometry_msgs::TwistStamped::ConstPtr& msg)
{
    local_vel = *msg;
}

void local_acc_target(const double ax, const double ay, const double az)
{
    mavros_msgs::PositionTarget pva_target;
    pva_target.coordinate_frame = 1;
    pva_target.header.stamp = ros::Time::now();

    pva_target.acceleration_or_force.x = ax;
    pva_target.acceleration_or_force.y = ay;
    pva_target.acceleration_or_force.z = az;

    pva_target.type_mask = pva_target.IGNORE_YAW + 
                pva_target.IGNORE_VX + pva_target.IGNORE_VY + pva_target.IGNORE_VZ+
                pva_target.IGNORE_PX + pva_target.IGNORE_PY + pva_target.IGNORE_PZ ;

    pva_target.yaw_rate = 0;

    local_pva_pub.publish(pva_target);
}

void local_vel_target(const double vx, const double vy, const double vz)
{
    mavros_msgs::PositionTarget pva_target;
    pva_target.coordinate_frame = 1;
    pva_target.header.stamp = ros::Time::now();

    pva_target.velocity.x = vx;
    pva_target.velocity.y = vy;
    pva_target.velocity.z = vz;

    pva_target.type_mask = pva_target.IGNORE_YAW + pva_target.FORCE + 
                pva_target.IGNORE_AFX + pva_target.IGNORE_AFY + pva_target.IGNORE_AFZ +
                pva_target.IGNORE_PX + pva_target.IGNORE_PY + pva_target.IGNORE_PZ ;

    pva_target.yaw_rate = 0;

    local_pva_pub.publish(pva_target);
}

void local_pva_target(const double px, const double py, const double pz,
        const double vx, const double vy, const double vz, const double ax,
        const double ay, const double az)
{
    mavros_msgs::PositionTarget pva_target;
    pva_target.coordinate_frame = 1;
    pva_target.header.stamp = ros::Time::now();

    pva_target.position.x = px;
    pva_target.position.y = py;
    pva_target.position.z = pz;
    pva_target.velocity.x = vx;
    pva_target.velocity.y = vy;
    pva_target.velocity.z = vz;
    pva_target.acceleration_or_force.x = ax;
    pva_target.acceleration_or_force.y = ay;
    pva_target.acceleration_or_force.z = az;

    pva_target.type_mask = pva_target.IGNORE_YAW;

    pva_target.yaw_rate = 0;

    local_pva_pub.publish(pva_target);
}

void takeoff(ros::NodeHandle &nh)
{
    

    //the setpoint publishing rate MUST be faster than 2Hz
    ros::Rate rate(20.0);

    // wait for FCU connection
    while(ros::ok() && !current_state.connected){
        ros::spinOnce();
        rate.sleep();
    }

    // geometry_msgs::PoseStamped pose;
    // pose.pose.position.x = 0;
    // pose.pose.position.y = 0;
    // pose.pose.position.z = 1.5;

    //send a few setpoints before starting
    for(int i = 100; ros::ok() && i > 0; --i){
        local_pos_pub.publish(ref_pose);
        ros::spinOnce();
        rate.sleep();
    }

    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";

    mavros_msgs::CommandBool arm_cmd;
    arm_cmd.request.value = true;

    ros::Time last_request = ros::Time::now();

    while((current_state.mode != "OFFBOARD") || (!arm_cmd.response.success))
    {
        if( current_state.mode != "OFFBOARD" &&(ros::Time::now() - last_request > ros::Duration(5.0)))
        {
            if( set_mode_client.call(offb_set_mode) &&offb_set_mode.response.mode_sent)
            {
                ROS_INFO("Offboard enabled");
            }
            last_request = ros::Time::now();
        }
        else 
        {
            if( !current_state.armed &&(ros::Time::now() - last_request > ros::Duration(5.0)))
            {
                if( arming_client.call(arm_cmd) &&arm_cmd.response.success)
                {
                    ROS_INFO("Vehicle armed");
                }
                last_request = ros::Time::now();
            }
        }
        local_pos_pub.publish(ref_pose);

        ros::spinOnce();
        rate.sleep();
    }

    double vx = - (local_pos.pose.position.x - ref_pose.pose.position.x);
    double vy = - (local_pos.pose.position.y - ref_pose.pose.position.y);
    double vz = - (local_pos.pose.position.z - ref_pose.pose.position.z);

    double dis = sqrt((local_pos.pose.position.x - ref_pose.pose.position.x)*(local_pos.pose.position.x - ref_pose.pose.position.x)
    + (local_pos.pose.position.y - ref_pose.pose.position.y)*(local_pos.pose.position.y - ref_pose.pose.position.y)
    + (local_pos.pose.position.z - ref_pose.pose.position.z)*(local_pos.pose.position.z - ref_pose.pose.position.z));
    while(dis > 0.1)
    {
        // local_pos_pub.publish(ref_pose);
        local_vel_target(vx, vy, vz);
        vx = - (local_pos.pose.position.x - ref_pose.pose.position.x);
        vy = - (local_pos.pose.position.y - ref_pose.pose.position.y);
        vz = - (local_pos.pose.position.z - ref_pose.pose.position.z);
        ros::spinOnce();
        rate.sleep();
        dis = sqrt((local_pos.pose.position.x - ref_pose.pose.position.x)*(local_pos.pose.position.x - ref_pose.pose.position.x)
            + (local_pos.pose.position.y - ref_pose.pose.position.y)*(local_pos.pose.position.y - ref_pose.pose.position.y)
            + (local_pos.pose.position.z - ref_pose.pose.position.z)*(local_pos.pose.position.z - ref_pose.pose.position.z));

    }


}

void trigger_cb(const std_msgs::Int8::ConstPtr &msg)
{
    if (msg->data == 1)
    {
        begin_task = 1;
    }
}

void move_to_init()
{
    ros::Rate rate(20.0);
    double vx = - (local_pos.pose.position.x - ref_pose.pose.position.x);
    double vy = - (local_pos.pose.position.y - ref_pose.pose.position.y);
    double vz = - (local_pos.pose.position.z - ref_pose.pose.position.z);

    double dis = sqrt((local_pos.pose.position.x - ref_pose.pose.position.x)*(local_pos.pose.position.x - ref_pose.pose.position.x)
    + (local_pos.pose.position.y - ref_pose.pose.position.y)*(local_pos.pose.position.y - ref_pose.pose.position.y)
    + (local_pos.pose.position.z - ref_pose.pose.position.z)*(local_pos.pose.position.z - ref_pose.pose.position.z));
    while(dis > 0.1)
    {
        // local_pos_pub.publish(ref_pose);
        local_vel_target(vx, vy, vz);
        vx = - (local_pos.pose.position.x - ref_pose.pose.position.x);
        vy = - (local_pos.pose.position.y - ref_pose.pose.position.y);
        vz = - (local_pos.pose.position.z - ref_pose.pose.position.z);
        ros::spinOnce();
        rate.sleep();
        dis = sqrt((local_pos.pose.position.x - ref_pose.pose.position.x)*(local_pos.pose.position.x - ref_pose.pose.position.x)
            + (local_pos.pose.position.y - ref_pose.pose.position.y)*(local_pos.pose.position.y - ref_pose.pose.position.y)
            + (local_pos.pose.position.z - ref_pose.pose.position.z)*(local_pos.pose.position.z - ref_pose.pose.position.z));

    }

}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "tube_plan_node");
    ros::NodeHandle nh;

    ros::Subscriber trigger_sub = nh.subscribe<std_msgs::Int8>("/tube/trigger", 10, trigger_cb);
    // state_sub = nh.subscribe<mavros_msgs::State>
    //         ("mavros/state", 10, state_cb);
    local_pos_sub = nh.subscribe<geometry_msgs::PoseStamped>
            ("mavros/local_position/pose", 10, local_pos_cb);
    local_vel_sub = nh.subscribe<geometry_msgs::TwistStamped>
            ("mavros/local_position/velocity_local", 10, local_vel_cb);        
    local_pos_pub = nh.advertise<geometry_msgs::PoseStamped>
            ("mavros/setpoint_position/local", 10);
    local_pva_pub = nh.advertise<mavros_msgs::PositionTarget>
            ("mavros/setpoint_raw/local",10);
    predict_pos_pub = nh.advertise<nav_msgs::Path>("mpc/perdict_pose", 10);
    ref_pos_pub = nh.advertise<nav_msgs::Path>("mpc/ref_pose", 10);
    // arming_client = nh.serviceClient<mavros_msgs::CommandBool>
    //         ("mavros/cmd/arming");
    // set_mode_client = nh.serviceClient<mavros_msgs::SetMode>
    //         ("mavros/set_mode");
    // trajectory planning
    Tube_planning::GeneratorFist generator_first;
    Tube_planning::polyCurve traj;
    
    generator_first.init(nh);
    traj = generator_first.trajPlan();
    std::cout << "trajectory planning success" << std::endl;
    // std::cout << "path: " << traj.Path << std::endl;
    // std::cout << "poly x: " << traj.Poly_x << std::endl;
    // std::cout << "poly y: " << traj.Poly_y << std::endl;
    // std::cout << "poly z: " << traj.Poly_z << std::endl;
    generator_first.evaluateTraj();


    // mpc 
    Tube_planning::mmpc MPC;
    Eigen::Array3d pos_init = MPC.init(nh, traj);
    ref_pose.pose.position.x = pos_init(0);
    ref_pose.pose.position.y = pos_init(1);
    ref_pose.pose.position.z = pos_init(2);
    std::cout << "set init pos: " << pos_init.transpose() << std::endl;


    // takeoff
    // std::cout << "begin takeoff" << std::endl;
    // takeoff(nh);
    // std::cout << "takeoff success" << std::endl; 
    
    ros::Rate rate(50);
    std::cout << "wait begin" << std::endl; 
    while (!begin_task)
    {
        // generator_first.evaluateTraj();
        ros::spinOnce();
        rate.sleep();
    }

    

    std::cout << "move to init pose" << std::endl; 
    move_to_init();

    MPC.setBeginTime(ros::Time::now().toSec());
    std::cout << "reset begin time " << std::endl;
    std::cout << " --------------- begin MPC ----------- " << std::endl;

    // MPC.getFeedback(local_pos, local_vel);
    // MPC.getReference();
    // Eigen::ArrayXd X = MPC.solver();

    while (ros::ok())
    {
        // local_pos_pub.publish(ref_pose);
        

        // for display traj
        // generator_first.evaluateTraj();

        MPC.getFeedback(local_pos, local_vel);
        Eigen::MatrixXd state_ref = MPC.getReference();
        nav_msgs::Path path_point_1;
        path_point_1.header.stamp = ros::Time::now();
        path_point_1.header.frame_id = "map";
        geometry_msgs::PoseStamped this_path_stamped_1;
        for (int i = 0; i < state_ref.cols(); i++)
        {
            /* code */
            this_path_stamped_1.header.stamp = ros::Time::now();
            this_path_stamped_1.header.frame_id = "map";
            this_path_stamped_1.pose.position.x = state_ref(0, i);
            this_path_stamped_1.pose.position.y = state_ref(1, i);
            this_path_stamped_1.pose.position.z = state_ref(2, i);

            path_point_1.poses.push_back(this_path_stamped_1);
        }
        ref_pos_pub.publish(path_point_1);


        Eigen::ArrayXd X = MPC.solver();
        // local_acc_target(X(6), X(7), X(8));
        if (MPC.getRuntime())
        {
            local_pva_target(X(2*9), X(2*9 + 1), X(2*9 + 2), 
                X(2*9 + 3), X(2*9 + 4), X(2*9 + 5),
                X(9 + 6), X(9 + 7), X(9 + 8));
        }
        else
        {
            // std::cout << "-----------------land--------------"<< std::endl;
            local_vel_target(0, 0, -0.5);
        }

       

        nav_msgs::Path path_point;
        path_point.header.stamp = ros::Time::now();
        path_point.header.frame_id = "map";
        geometry_msgs::PoseStamped this_path_stamped;
        for (int i = 0; i < 40; i++)
        {
            /* code */
            this_path_stamped.header.stamp = ros::Time::now();
            this_path_stamped.header.frame_id = "map";
            this_path_stamped.pose.position.x = X(i*9);
            this_path_stamped.pose.position.y = X(i*9+1);
            this_path_stamped.pose.position.z = X(i*9+2);

            path_point.poses.push_back(this_path_stamped);
        }
        predict_pos_pub.publish(path_point);


        ros::spinOnce();
        rate.sleep();
    }
    std::cout << "-----------------land--------------"<< std::endl;

    
    return 0;
}
